import numpy as np
import cv2

def sort_quad_vertices(vertices):
    """
    功能：将四边形顶点按「左上角→右上角→右下角→左下角」顺序排序
    输入：vertices - 多边形逼近得到的顶点数组（shape=(4,1,2)）
    输出：sorted_pts - 排序后的顶点数组（shape=(4,2)，float32）
    """
    # 去除冗余维度，转为4行2列的坐标数组
    pts = vertices.reshape((4, 2)).astype(np.float32)

    # 1. 按x+y之和排序：左上角（sum最小）、右下角（sum最大）
    sum_xy = pts.sum(axis=1)
    tl_idx = np.argmin(sum_xy)  # 左上角索引
    br_idx = np.argmax(sum_xy)  # 右下角索引

    # 2. 排除已确定的角点，剩余两个为右上角和左下角
    remaining_idx = [i for i in range(4) if i not in [tl_idx, br_idx]]
    tr_idx, bl_idx = remaining_idx

    # 3. 区分右上角和左下角：x坐标大的是右上角
    if pts[tr_idx, 0] < pts[bl_idx, 0]:
        tr_idx, bl_idx = bl_idx, tr_idx

    # 按顺序返回四个角点
    sorted_pts = np.array([pts[tl_idx], pts[tr_idx], pts[br_idx], pts[bl_idx]], dtype=np.float32)
    return sorted_pts


def get_target_dimensions(sorted_pts):
    """
    功能：根据排序后的顶点计算目标图像真实宽高（避免拉伸）
    输入：sorted_pts - 排序后的顶点数组（shape=(4,2)）
    输出：target_w, target_h - 目标宽高（整数）
    """
    tl, tr, br, bl = sorted_pts

    # 计算水平方向宽度（取上下边最大值，确保完整覆盖）
    width_top = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    width_bottom = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    target_w = int(np.ceil(max(width_top, width_bottom)))  # 向上取整避免裁剪

    # 计算垂直方向高度（取左右边最大值）
    height_left = np.sqrt(((bl[0] - tl[0]) ** 2) + ((bl[1] - tl[1]) ** 2))
    height_right = np.sqrt(((br[0] - tr[0]) ** 2) + ((br[1] - tr[1]) ** 2))
    target_h = int(np.ceil(max(height_left, height_right)))

    return target_w, target_h


def image_correction(file_path):
    img = cv2.imread(file_path)

    # 灰度化
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 二值化
    ret, img_binary = cv2.threshold(img_gray, 127, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)

    # 高斯滤波
    img_gaussian = cv2.GaussianBlur(img_binary, ksize=(5, 5), sigmaX=3)

    # 寻找轮廓
    contours, hierarchy = cv2.findContours(
        img_gaussian,
        mode=cv2.RETR_LIST,   # 轮廓检索模式
        method=cv2.CHAIN_APPROX_SIMPLE  # 轮廓 Approximation 模式
    )

    print("轮廓数量：", len(contours))
    for i, c in enumerate(contours):
        print(f"第{i+1}个轮廓的边界点的个位数为{len(c)}")
    # 获取最大轮廓
    max_contour = max(contours, key=cv2.contourArea)
    print(len(max_contour))

    img_copy = img.copy()
    # 多边形逼近
    perimeter = cv2.arcLength(max_contour, True)
    approx = cv2.approxPolyDP(max_contour, 0.04 * perimeter, True)
    cv2.drawContours(img_copy, [approx], -1, (0, 255, 0), 2)
    print(approx)
    # cv2.drawContours(img, [max_contour], -1, (0, 0, 255), 2)

    # 顶点排序（关键：确保与目标顶点顺序对应）
    sorted_src_pts = sort_quad_vertices(approx)
    # 计算目标图像尺寸
    target_w, target_h = get_target_dimensions(sorted_src_pts)
    print(f"矫正后目标尺寸：宽={target_w}px，高={target_h}px")

    # 定义目标顶点（与源顶点顺序严格一致：左上→右上→右下→左下）
    dst_pts = np.float32([
        [0, 0],  # 目标左上
        [target_w - 1, 0],  # 目标右上（减1避免坐标越界）
        [target_w - 1, target_h - 1],  # 目标右下
        [0, target_h - 1]  # 目标左下
    ])

    # 透视变换
    M = cv2.getPerspectiveTransform(sorted_src_pts, dst_pts)
    img_warp = cv2.warpPerspective(img, M, (target_w, target_h))
    return img_warp



